/*
 * Copyright 2015 MingxingWang

 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at

 *     http://www.apache.org/licenses/LICENSE-2.0

 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.spark.graphx.graphdb

import org.apache.spark.graphx.{Edge, VertexId}
import org.apache.spark.graphx.util.collection.GraphXPrimitiveKeyOpenHashMap
import org.apache.spark.util.collection.BitSet

import scala.collection.mutable
import scala.collection.mutable.HashSet
import scala.reflect.ClassTag

/**
 * Created by Ming on 2015/3/26.
 */
class LEdgePartition[@specialized(Char, Int, Boolean, Byte, Long, Float, Double) ED: ClassTag, VD: ClassTag]
(
  val srcIds: Array[VertexId],
  val dstIds: Array[VertexId],
  val data: Array[ED],
  val index: GraphXPrimitiveKeyOpenHashMap[VertexId, Int]
  ){
  val mask: BitSet = new BitSet(srcIds.size)

  val size: Int = srcIds.size - mask.cardinality()

  def isDefined(srcId: VertexId, dstId: VertexId): Boolean = {
    var i = index.getOrElse(srcId, -1)
    if (i >= 0) {
      while (i < srcIds.length && srcIds(i) == srcId && dstId <= dstIds(i)) {
        if (dstId == dstIds(i) && !mask.get(i)) {
          return true
        }
        i = i + 1
      }
    }
    false
  }

  def delete(srcId: VertexId, dstId: VertexId): Boolean = {
    var i = index.getOrElse(srcId, -1)
    if (i >= 0) {
      while (i < srcIds.length && srcIds(i) == srcId && dstId <= dstIds(i)) {
        if (dstId == dstIds(i) && !mask.get(i)) {
          mask.set(i)
          return true
        }
        i = i + 1
      }
    }
    false
  }

  def find(srcId: VertexId, dstId: VertexId): Option[ED] = {
    var i = index.getOrElse(srcId, -1)
    if (i >= 0) {
      while (i < srcIds.length && srcIds(i) == srcId && dstId <= dstIds(i)) {
        if (dstId == dstIds(i) && !mask.get(i)) {
          return Some(data(i))
        }
        i = i + 1
      }
    }
    None
  }

  def update(srcId: VertexId, dstId: VertexId, vd: ED): Boolean = {
    var i = index.getOrElse(srcId, -1)
    if (i >= 0) {
      while (i < srcIds.length && srcIds(i) == srcId && dstId <= dstIds(i)) {
        if (dstId == dstIds(i) && !mask.get(i)) {
          data(i) = vd
          return true
        }
        i = i + 1
      }
    }
    false
  }

  def findOutEdges(srcId: VertexId,
                   filter: ED => Boolean = ED => true)
  : Map[VertexId, ED] = {
    val es = new mutable.HashMap[VertexId, ED]()
    try {
      var i = index.getOrElse(srcId, -1)
      if (i >= 0) {
        while (i < srcIds.length && srcIds(i) == srcId) {
          if (!mask.get(i) && filter(data(i))) {
            es += (dstIds(i) -> data(i))
          }
          i = i + 1
        }
      }
    } catch {
      case _: Throwable => {
        println("error in findOutEdges")
      }
    }
    es.toMap
  }

  def findOutVertices(srcId: VertexId,
                      filter: ED => Boolean = ED => true)
  : Set[VertexId] = {
    val es = new HashSet[VertexId]()
    var i = index.getOrElse(srcId, -1)
    if (i >= 0) {
      while (i < srcIds.length && srcIds(i) == srcId) {
        if (!mask.get(i) && filter(data(i))) {
          es += dstIds(i)
        }
        i = i + 1
      }
    }
    es.toSet
  }

  /**
   * Get an iterator over the edges in this partition.
   *
   * Be careful not to keep references to the objects from this iterator.
   * To improve GC performance the same object is re-used in `next()`.
   *
   * @return an iterator over edges in the partition
   */
  def iterator = new Iterator[Edge[ED]] {
    private[this] val edge = new Edge[ED]
    private[this] var pos = 0

    override def hasNext: Boolean = {
      while(pos < LEdgePartition.this.size && mask.get(pos)){
        pos += 1
      }
      pos < LEdgePartition.this.size
    }

    override def next(): Edge[ED] = {
      edge.srcId = srcIds(pos)
      edge.dstId = dstIds(pos)
      edge.attr = data(pos)
      pos += 1
      edge
    }
  }

  /**
   * Reverse all the edges in this partition.
   *
   * @return a new edge partition with all edges reversed.
   */
  def reverse: LEdgePartition[ED, VD] = {
    val builder = new LEdgePartitionBuilder[ED, VD](size)
    for (e <- iterator) {
      builder.add(e.dstId, e.srcId, e.attr)
    }
    builder.toEdgePartition
  }
}
